import os
import time
from llama_index.core.schema import Document
from llama_index.core.response.pprint_utils import pprint_source_node
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex


def load_by_line(data_dir):
    texts = []
    for file in os.listdir(data_dir):
        for line in open(f"{data_dir}/{file}", "r", encoding="utf8"):
            line = line.strip()
            texts.append(line)
    return texts

# setting embeding model
Settings.embed_model = "local:/Users/weigang/D/LLM-Model/m3e-base"   # 本地模型需要加"local:"开头
Settings.chunk_size = 512                         # chunk_size大小

# load data
#data_path = '/llm/data'
#documents = SimpleDirectoryReader(data_path).load_data()

#texts = load_by_line(data_path)
texts = ["周杰伦", "蔡徐坤", "王捷", "王杰", "冰箱", "空调", "小轿车", "飞机", "西游记", "红楼梦", "三国演义", "水浒传"]
documents = [Document(text=txt) for txt in texts]

# build index
index = VectorStoreIndex.from_documents(documents)

# 将embedding后的结果保存到本地，这样下次使用直接加载就不用重复操作。
from llama_index.core import StorageContext, load_index_from_storage
# Persisting to disk
index.storage_context.persist(persist_dir="<persist_dir>")
# rebuild storage context
storage_context = StorageContext.from_defaults(
        docstore=SimpleDocumentStore.from_persist_dir(persist_dir="<persist_dir>"),
        index_store=SimpleIndexStore.from_persist_dir(persist_dir="<persist_dir>"),
        persist_dir="<persist_dir>")
storage_context = StorageContext.from_defaults(persist_dir="<persist_dir>")
# load index
index = load_index_from_storage(storage_context)


# as retriever
ES = index.as_retriever(similarity_top_k=5)
while True:
    print("\n")
    try:
        query = input("query:")
        print("=============================================")
    except Exception as e:
        print(e)
        continue
    if query == "exit":
        break
    start_time = time.time()
    res = ES.retrieve(query)
    cost_time = time.time() - start_time
    print(f"cost time: {cost_time*1000}ms")
    for idx, result in enumerate(res):
        pprint_source_node(result)